import time
import os

from langchain.tools import tool
from langchain.llms import Ollama
from crewai import Agent, Task, Process, Crew
from langchain_openai import ChatOpenAI
from jwt_token import generate_token
import time

def get_glm(temprature):
    llm = ChatOpenAI(
        model_name="glm-4",
        openai_api_base="https://open.bigmodel.cn/api/paas/v4",
        openai_api_key=generate_token(os.environ["ZHIPU_API_KEY"], 3600),
        streaming=False,
        temperature=temprature
    )
    return llm


from langchain.agents import load_tools

# To load Human in the loop
human_tools = load_tools(["human"])


# To Load Local models through Ollama
# gemma = Ollama(model="gemma:7b")
glm = get_glm(0.5)

class RateLimiter:
    def __init__(self, min_interval_seconds):
        self.min_interval_seconds = min_interval_seconds
        self.last_request_time = None

    def wait_for_next_request(self):
        if self.last_request_time is not None:
            elapsed_time = time.time() - self.last_request_time
            wait_time = self.min_interval_seconds - elapsed_time
            if wait_time > 0:
                time.sleep(wait_time)
        self.last_request_time = time.time()

@tool
def rate_limited_request(input_str: str, rate_limiter: RateLimiter) -> str:
    """This tool wraps the actual API request, ensuring it respects the rate limit.
    
    Args:
        input_str: The input string for the API request.
        rate_limiter: An instance of the RateLimiter class to enforce rate limits.
        
    Returns:
        The result from the API call.
    """
    # Wait for the next permissible request time if necessary.
    rate_limiter.wait_for_next_request()
    
    # Here, you would make the actual API request.
    # For demonstration purposes, we'll just return the input string.
    return input_str

agnent_exporter = Agent(
    role="你是一个selenium爬虫专家",
    goal="你能够根据用户提供的网址分析其中的内容",
    backstory="""你是一位擅长规划多智能体的专家策略家。 
    你擅长把用户的需求，一步一步的分析为可行的多智能体设计方案。
    """,
    verbose=True,
    allow_delegation=False,
    llm=glm,  # 移除以使用默认的gpt-4
    tools=[rate_limited_request, RateLimiter]
)

programmer = Agent(
    role="CrewAI开发工程师",
    goal="使用CrewAI的语法，完成应用的开发",
    backstory="""您是 CrewAI 应用开发方面的专家，擅长把整个多智能体应用的代码开发完成。""",
    verbose=True,
    allow_delegation=True,
    llm=glm,  # remove to use default gpt-4
    tools=[rate_limited_request, RateLimiter]
)

task_report = Task(
    description="""生成一个知乎问答专家的多智能体应用，该应用可以根据用户输入的知乎问题，生成一篇有深度的文章，用来回答该知乎问题。文章中要有案例，有故事，易于理解，没有废话。中间不需要利用知识图谱等深层技术进行复杂提取，仅用crewAI能实现的方法即可，请用中文。
    """,
    agent=agnent_exporter,
)

program = Task(
    description="""使用CrewAI写一个能够为知乎问题生成高赞回答的应用，该应用可以根据用户输入的知乎问题，生成一篇有深度的文章，用来回答该知乎问题。文章中要有案例，有故事，易于理解，没有废话。中间不需要利用知识图谱等深层技术进行复杂提取，仅用crewAI能实现的方法即可，请用中文。
    """,
    agent=programmer,
)

# instantiate crew of agents
crew = Crew(
    agents=[agnent_exporter, programmer],
    tasks=[task_report, program],
    verbose=2,
    process=Process.sequential,  # 顺序处理意味着任务将按顺序执行，每个任务执行后，其结果将被作为附加内容传递给下一个任务。
)

# Get your crew to work!
result = crew.kickoff()

print("######################")
print(result)